Hi @Lilin-Huang. SHy is a non-trivial model to implement (hypergraph message passing, Gumbel-Softmax phenotype extraction, multi-objective loss), and you have it all in working form with real ablation numbers from MIMIC-IV. The BaseModel and BaseTask integration is correct, the __init__.py changes are purely additive, and the ablation tables in the PR description are the kind of honest reproducibility reporting we want to see. A few things worth addressing.

Scalability: forward is sample-by-sample.

for i in range(batch_size):
    H = self._build_incidence_matrix(codes_batch[i]).to(self.device)
    ...
    tp_mats, tp_embs = self._encode_patient(X, H)

Each sample gets its own incidence matrix, HGNN pass, K extractors, and decoder call inside a Python loop. This will not scale beyond the 1000-patient dev subset used in your ablations. Per-patient hypergraphs make full batching nontrivial, but block-diagonal stacking (a standard trick from PyTorch Geometric) would let you batch the HGNN and phenotype extraction.

Not a merge blocker, but please add a note in the class docstring documenting the expected scalability (e.g., "This implementation processes samples sequentially in forward. For large batches or datasets, consider batching via block-diagonal hypergraph construction").

logit in the output dict is actually probabilities.

pred = torch.sigmoid(self.predict(combined))
...
return {"loss": loss, "y_prob": pred, "y_true": y_true, "logit": pred}

pred is post-sigmoid. Other PyHealth models return pre-sigmoid logits in the logit key. This will confuse downstream code that expects to apply sigmoid itself, for example calibration or temperature scaling. Could you either return the pre-sigmoid linear output in logit, or document in the docstring that logit is probability-valued for this model?

Use ValueError instead of assert for argument validation.

assert len(self.label_keys) == 1, "SHy supports exactly one label key (multilabel)"
assert len(self.feature_keys) == 1, "SHy expects exactly one feature key (nested_sequence)"

Asserts are stripped when Python is run with -O. For argument validation that should always run, raise ValueError instead. Small fix, two lines each.

Empty-history samples contribute to prediction loss from zero-vector input.

if H.sum() == 0:
    zero = torch.zeros(self.num_tp, self.hidden_dim, device=self.device)
    latent_list.append(zero)
    continue

When a patient has an all-zero incidence matrix (possible if all codes got clipped or visits are empty), the classifier still runs on the zero vector and the loss still counts. The task filters to len(visits) >= 2, so this should be rare in practice, but if it does happen the model learns to predict whatever the dataset average is for that position. Could you either skip these samples in the loss (exclude their indices from prediction BCE), or add a comment explaining the trade-off?

Fidelity loss does not handle class imbalance.

Your prediction loss correctly reweights positives:

num_pos = y_true.sum(dim=1, keepdim=True).clamp(min=1)
num_neg = (y_true.shape[1] - num_pos).clamp(min=1)
pos_weight = (num_neg / num_pos).expand_as(y_true)

But the fidelity loss is plain BCE on the incidence matrix:

F.binary_cross_entropy(r.clamp(1e-9, 1 - 1e-9), h.float())

The incidence matrix is typically very sparse (most codes are not in a visit). Without reweighting, fidelity encourages predicting all zeros, which is the uninformative solution. Your ablation shows fidelity has small positive weight (0.1) so this may be contained, but worth applying the same pos_weight treatment or commenting on why raw BCE is intended.

Tests verify it runs but do not verify the core model behavior.

The 12 tests cover shapes, output keys, backward, probability range, and a few hyperparameter variations. None of them check:

1. That for num_tp=K, the model actually produces K distinct phenotype sub-hypergraphs (e.g., len(tp_list[0]) == K inside the forward path).
2. That different configurations produce different outputs (for example, that num_tp=1 and num_tp=5 give measurably different losses on the same input).
3. That the add_ratio false-negative addition changes the incidence matrix in a non-trivial way.

These would strengthen the PR by making it harder for future refactors to silently break the paper's core mechanism. A test that inspects the shape and sparsity of what PhenotypeExtractor produces would be a good start.

Teacher forcing in the decoder during evaluation.

# Teacher forcing
prev_codes = H.T[t]

The decoder uses ground-truth codes from visit t as input for predicting visit t+1, both during training and evaluation. During eval this means the fidelity reconstruction is computed with ground-truth inputs, which inflates the reconstruction quality. For the paper's reported fidelity numbers this is fine (they are training-time regularization metrics), but if anyone looks at eval fidelity they will over-interpret it. Could you gate teacher forcing on self.training, or add a docstring note?

Smaller items

1. The docstring example for SHy:

>>> out = model(**batch)
>>> out["loss"]

Shows no expected output. Consider either removing the last line or adding a torch.Size(...) or tensor(...) to make it a valid doctest.

2. _build_incidence_matrix uses a Python for-loop over visits. scatter_ on a flattened index tensor would vectorize this. Minor perf.

3. num_edges = E.max().item() + 1 in UniGINConv.forward crashes if E is empty. The outer if H.sum() == 0: continue in forward guards against this, but a defensive check in the layer would make it safer for direct use.

4. The ablation table in the PR is valuable but would be more useful if you also pasted it (or a reduced version) into the example script's top docstring, so users can see the expected output without running the full MIMIC-IV pipeline.

5. "pred only" in the ablation table is confusing as a label. Could you rename to "no auxiliary loss" or similar for clarity?